from typing import Optional, Tuple, Union

import torch
from torch import Tensor
from mmdeploy.core import SYMBOLIC_REWRITER, FUNCTION_REWRITER
from mmcv.ops.modulated_deform_conv import ModulatedDeformConv2dFunction


@SYMBOLIC_REWRITER.register_symbolic(
    'mmcv.ops.modulated_deform_conv.ModulatedDeformConv2dFunction',
    backend='ascend')
def modulated_deformconv2dfunction_symbolic_ascend(
        g, input, offset, mask, weight, bias,
        stride, padding, dilation, groups, deform_groups
    ):
    input_tensors = [input, weight, offset]
    if bias is not None:
        input_tensors.append(bias)
    return g.op('mmdeploy::DeformableConv2D',
                *input_tensors,
                strides_i=stride,
                pads_i=padding,
                dilations_i=dilation,
                groups_i=groups,
                data_format_s='NCHW',
                deformable_groups_i=deform_groups)


@FUNCTION_REWRITER.register_rewriter(
    'mmcv.ops.modulated_deform_conv.modulated_deform_conv2d',
    backend='ascend')
def modulated_deform_conv__ascend(
        input, offset, mask, weight, bias,
        stride, padding, dilation, groups, deform_groups
    ):
    """rewriter for the custom ascend mdcn op."""
    offset_y = offset.reshape(1, -1, 2, offset.shape[2],
        offset.shape[3])[:, :, 0, ...].reshape(1, offset.shape[1] // 2,
        offset.shape[2], offset.shape[3])
    offset_x = offset.reshape(1, -1, 2, offset.shape[2],
        offset.shape[3])[:, :, 1, ...].reshape(1, offset.shape[1] // 2,
        offset.shape[2], offset.shape[3])
    offset = torch.cat((offset_x, offset_y, mask), 1)

    return ModulatedDeformConv2dFunction.apply(
        input, offset, mask, weight, bias,
        stride, padding, dilation, groups, deform_groups
    )


# @FUNCTION_REWRITER.register_rewriter(
#     'mmcv.ops.modulated_deform_conv.ModulatedDeformConv2dPack.forward',
#     backend='ascend')
# def pack_forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
#     out = self.conv_offset(x)
#     o1, o2, mask = torch.chunk(out, 3, dim=1)
#     offset = torch.cat((o1, o2), dim=1)
#     mask = torch.sigmoid(mask)
#     if torch.onnx.is_in_onnx_export():
#         offset_y = offset.reshape(1, -1, 2, offset.shape[2].numpy(),
#                 offset.shape[3].numpy())[:, :, 0, ...].reshape(1,
#                 offset.shape[1].numpy() // 2, offset.shape[2].numpy(),
#                 offset.shape[3].numpy())
#         offset_x = offset.reshape(1, -1, 2, offset.shape[2].numpy(),
#                 offset.shape[3].numpy())[:, :, 1, ...].reshape(1,
#                 offset.shape[1].numpy() // 2, offset.shape[2].numpy(),
#                 offset.shape[3].numpy())
#         offset = torch.cat((offset_x, offset_y, mask), 1)
#     return ModulatedDeformConv2dFunction.apply(
#         x, offset, mask, self.weight, self.bias,
#                                     self.stride, self.padding,
#                                     self.dilation, self.groups,
#                                     self.deform_groups)
